#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
基于向量搜索的树形结构Neo4j查询API服务
基于FastAPI提供Web接口，使用向量相似性搜索
"""

from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional, Dict, Any
import uvicorn
from vector_tree_query_tool import VectorTreeQueryTool
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 创建FastAPI应用
app = FastAPI(
    title="基于向量搜索的树形结构Neo4j查询API",
    description="提供决策树知识库的向量相似性查询接口",
    version="2.0.0"
)

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 生产环境中应该限制具体域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 全局查询工具实例
query_tool = None

# 请求模型
class QueryRequest(BaseModel):
    query: str
    limit: Optional[int] = None

class NodeQueryRequest(BaseModel):
    node_name: str
    product_code: Optional[str] = None

# 响应模型
class NodeResult(BaseModel):
    product_code: str
    node_type: str
    content: str
    filename: str
    line_number: int
    relevance_score: float
    vector_similarity: Optional[float] = None

class ChainStep(BaseModel):
    content: str
    node_type: str
    is_matched: bool

class ChainResult(BaseModel):
    product_code: str
    filename: str
    line_number: int
    chain: List[ChainStep]

class QueryResponse(BaseModel):
    success: bool
    message: str
    result_type: str  # "nodes" 或 "chains"
    total_count: int
    nodes: Optional[List[NodeResult]] = None
    chains: Optional[List[ChainResult]] = None

# 启动事件
@app.on_event("startup")
async def startup_event():
    """应用启动时初始化查询工具"""
    global query_tool
    try:
        query_tool = VectorTreeQueryTool()
        if query_tool.neo4j_driver and query_tool.milvus_collection:
            logger.info("✅ 向量搜索查询工具初始化成功")
        else:
            logger.error("❌ 向量搜索查询工具初始化失败")
            raise Exception("无法连接到Neo4j或Milvus数据库")
    except Exception as e:
        logger.error(f"❌ 启动失败: {e}")
        raise

# 关闭事件
@app.on_event("shutdown")
async def shutdown_event():
    """应用关闭时清理资源"""
    global query_tool
    if query_tool:
        query_tool.close()
        logger.info("✅ 数据库连接已关闭")

def search_by_node_name_logic(query_tool, node_name, product_code=None):
    """
    根据节点名称查询其子节点的核心逻辑
    
    Args:
        query_tool: 查询工具实例
        node_name: 父节点名称
        product_code: 可选的产品编号
    
    Returns:
        list: 子节点查询结果列表
    """
    return query_tool.search_by_node_name(node_name, product_code)

@app.post("/search", response_model=QueryResponse)
async def search_query(request: QueryRequest):
    """
    向量搜索查询接口
    
    Args:
        request: 查询请求，包含查询文本和可选的限制数量
    
    Returns:
        QueryResponse: 查询结果
    """
    global query_tool
    
    if not query_tool or not query_tool.neo4j_driver or not query_tool.milvus_collection:
        raise HTTPException(status_code=503, detail="数据库连接不可用")
    
    try:
        # 执行向量搜索查询
        results = query_tool.search_tree_paths_with_vector(request.query, request.limit)
        
        if not results:
            return QueryResponse(
                success=True,
                message="没有找到匹配的结果",
                result_type="nodes",
                total_count=0,
                nodes=[]
            )
        
        # 检查结果类型
        if 'chain' in results[0]:
            # 链路结果
            chains = []
            for result in results:
                chain_steps = []
                for step in result['chain']:
                    chain_steps.append(ChainStep(
                        content=step['content'],
                        node_type=step['node_type'],
                        is_matched=step.get('is_matched', False)
                    ))
                
                chains.append(ChainResult(
                    product_code=result['product_code'],
                    filename=result['filename'],
                    line_number=result['line_number'],
                    chain=chain_steps
                ))
            
            return QueryResponse(
                success=True,
                message=f"找到 {len(chains)} 条决策链路",
                result_type="chains",
                total_count=len(chains),
                chains=chains
            )
        else:
            # 节点结果
            nodes = []
            for result in results:
                nodes.append(NodeResult(
                    product_code=result['product_code'],
                    node_type=result['node_type'],
                    content=result['content'],
                    filename=result['filename'],
                    line_number=result['line_number'],
                    relevance_score=result['relevance_score'],
                    vector_similarity=result.get('vector_similarity')
                ))
            
            return QueryResponse(
                success=True,
                message=f"找到 {len(nodes)} 个相关节点",
                result_type="nodes",
                total_count=len(nodes),
                nodes=nodes
            )
    
    except Exception as e:
        logger.error(f"向量搜索查询执行失败: {e}")
        raise HTTPException(status_code=500, detail=f"查询执行失败: {str(e)}")

@app.get("/search", response_model=QueryResponse)
async def search_query_get(
    q: str = Query(..., description="查询文本"),
    limit: Optional[int] = Query(None, description="结果数量限制")
):
    """
    GET方式的向量搜索查询接口
    
    Args:
        q: 查询文本
        limit: 可选的结果数量限制
    
    Returns:
        QueryResponse: 查询结果
    """
    request = QueryRequest(query=q, limit=limit)
    return await search_query(request)

@app.post("/search_node", response_model=QueryResponse)
async def search_by_node_name(request: NodeQueryRequest):
    """
    根据节点名称查询接口
    
    Args:
        request: 节点查询请求，包含节点名称和可选的产品编号
    
    Returns:
        QueryResponse: 查询结果
    """
    global query_tool
    
    if not query_tool or not query_tool.neo4j_driver:
        raise HTTPException(status_code=503, detail="数据库连接不可用")
    
    try:
        # 执行节点名称查询
        results = search_by_node_name_logic(query_tool, request.node_name, request.product_code)
        
        if not results:
            return QueryResponse(
                success=True,
                message="没有找到匹配的节点",
                result_type="nodes",
                total_count=0,
                nodes=[]
            )
        
        # 检查结果类型
        if 'chain' in results[0]:
            # 链路结果
            chains = []
            for result in results:
                chain_steps = []
                for step in result['chain']:
                    chain_steps.append(ChainStep(
                        content=step['content'],
                        node_type=step['node_type'],
                        is_matched=step.get('is_matched', False)
                    ))
                
                chains.append(ChainResult(
                    product_code=result['product_code'],
                    filename=result['filename'],
                    line_number=result['line_number'],
                    chain=chain_steps
                ))
            
            return QueryResponse(
                success=True,
                message=f"找到 {len(chains)} 条决策链路",
                result_type="chains",
                total_count=len(chains),
                chains=chains
            )
        else:
            # 节点结果
            nodes = []
            for result in results:
                nodes.append(NodeResult(
                    product_code=result['product_code'],
                    node_type=result['node_type'],
                    content=result['content'],
                    filename=result['filename'],
                    line_number=result['line_number'],
                    relevance_score=result['relevance_score'],
                    vector_similarity=result.get('vector_similarity')
                ))
            
            return QueryResponse(
                success=True,
                message=f"找到 {len(nodes)} 个相关节点",
                result_type="nodes",
                total_count=len(nodes),
                nodes=nodes
            )
    
    except Exception as e:
        logger.error(f"节点查询执行失败: {e}")
        raise HTTPException(status_code=500, detail=f"节点查询执行失败: {str(e)}")

@app.get("/search_node", response_model=QueryResponse)
async def search_by_node_name_get(
    node_name: str = Query(..., description="节点名称"),
    product_code: Optional[str] = Query(None, description="产品编号")
):
    """
    GET方式的节点名称查询接口
    
    Args:
        node_name: 节点名称
        product_code: 可选的产品编号
    
    Returns:
        QueryResponse: 查询结果
    """
    request = NodeQueryRequest(node_name=node_name, product_code=product_code)
    return await search_by_node_name(request)

@app.get("/query", response_model=QueryResponse)
async def query_get(
        q: str = Query(..., description="查询文本"),
        limit: Optional[int] = Query(None, description="结果数量限制")
):
    """
    GET方式的查询接口（兼容性接口）

    Args:
        q: 查询文本
        limit: 可选的结果数量限制

    Returns:
        QueryResponse: 查询结果
    """
    request = QueryRequest(query=q, limit=limit)
    return await search_query(request)

@app.get("/health")
async def health_check():
    """健康检查接口"""
    global query_tool

    status = {
        "status": "healthy",
        "neo4j_connected": False,
        "milvus_connected": False,
        "version": "2.0.0"
    }

    if query_tool:
        status["neo4j_connected"] = query_tool.neo4j_driver is not None
        status["milvus_connected"] = query_tool.milvus_collection is not None

    if not status["neo4j_connected"] or not status["milvus_connected"]:
        status["status"] = "unhealthy"
        raise HTTPException(status_code=503, detail=status)

    return status

@app.get("/")
async def root():
    """根路径接口"""
    return {
        "message": "基于向量搜索的树形结构Neo4j查询API",
        "version": "2.0.0",
        "features": [
            "向量相似性搜索",
            "语义理解查询",
            "决策树路径查询",
            "混合搜索策略"
        ],
        "endpoints": {
            "search": "POST/GET /search - 向量搜索查询",
            "search_node": "POST/GET /search_node - 节点名称查询",
            "query": "GET /query - 兼容性查询接口",
            "health": "GET /health - 健康检查"
        }
    }

if __name__ == "__main__":
    # 运行服务器
    uvicorn.run(
        "vector_tree_query_api:app",
        host="0.0.0.0",
        port=8001,  # 使用不同的端口避免冲突
        reload=True,
        log_level="info"
    )
